from typing import Any, Optional, Union
from haystack import Document, component, Pipeline
from haystack.tools import ComponentTool
from haystack.components.websearch import SerperDevWebSearch
from haystack.utils import Secret
from haystack.components.tools.tool_invoker import ToolInvoker
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.dataclasses import ChatMessage

# Create a SerperDev search component




@component
class WebSearch:
   
    def __init__(
        self
    ):
        pass
         
    

    @component.output_types(documents=list[Document], links=list[str])
    def run(self, query: str) -> dict[str, Union[list[Document], list[str]]]:
        return {"documents": [Document()], "links": ['ok']}
        
        
search = WebSearch()

# Create a tool from the component
tool = ComponentTool(
    component=search,
    name="web_search",  # Optional: defaults to "serper_dev_web_search"
    description="Search the web for current information on any topic"  # Optional: defaults to component docstring
)

# Create pipeline with OpenAIChatGenerator and ToolInvoker
chat_generator=OpenAIChatGenerator(model="DeepSeek-R1",api_base_url='https://api.deepseek.com',api_key='sk-605e60a1301040759a821b6b677556fb',tools=[tool])
pipeline = Pipeline()
pipeline.add_component("llm", chat_generator)
pipeline.add_component("tool_invoker", ToolInvoker(tools=[tool]))

# Connect components
pipeline.connect("llm.replies", "tool_invoker.messages")

message = ChatMessage.from_user("Use the web search tool to find information about Nikola Tesla")

# Run pipeline
result = pipeline.run({"llm": {"messages": [message]}})

print(result)